package com.hex.tcpproxy.util;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;

import java.net.InetSocketAddress;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author: [wujie]
 * @createTime: [2019/10/8 14:38]
 * @description:
 */
public abstract class ProxyHandler extends ChannelInboundHandlerAdapter {
    private Map<Object, ChannelHandlerContext> map = new ConcurrentHashMap<>();

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        tryRegister(ctx);
    }


    /**
     * 收到消息，自动触发
     *
     * @param ctx
     * @param msg
     * @throws Exception
     */
    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        String key = tryRegister(ctx);
        ByteBuf buf = (ByteBuf) msg;
        byte[] bytes = new byte[buf.readableBytes()];
        buf.readBytes(bytes);
        read(key, bytes);
    }

    public String tryRegister(ChannelHandlerContext ctx) {
        InetSocketAddress insocket = (InetSocketAddress) ctx.channel().remoteAddress();
        String key = insocket.getAddress().getHostAddress() + ":" + insocket.getPort();
        if (!map.keySet().contains(key)) {
            map.put(key, ctx);
        }
        return key;
    }

    public void write(String address, byte[] bytes) {
        map.get(address).writeAndFlush(Unpooled.copiedBuffer(bytes));
    }

    public void writeAll(byte[] bytes) {
        if(bytes[0]==0){
            StackTraceElement[] stackTrace = Thread.currentThread().getStackTrace();
            System.out.println();
        }
        map.forEach((k, v) -> {
            v.writeAndFlush(Unpooled.copiedBuffer(Unpooled.copiedBuffer(bytes)));
        });
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        cause.printStackTrace();
        String key = tryRegister(ctx);
        map.remove(key);
        /**当发生异常时，关闭 ChannelHandlerContext，释放和它相关联的句柄等资源 */
        ctx.close();
    }


    public abstract void read(String address, byte[] bytes);


}
